// plain markov chain text generation

use std::fs::File;
use std::io;
use std::io::{BufRead, BufReader};
use std::{collections::HashMap, hash::Hash, io::Error};
use std::{fmt::Debug, sync::Arc};

use once_cell::sync::Lazy;

extern crate sliding_windows;
use sliding_windows::{IterExt, Storage};

use futures::future::ok;

use crate::{DiscreteDistribution, common::NextTokenDistributionProvider};

extern crate regex;
use regex::Regex;

// #[derive(Clone)]
pub struct TokenIterator {
    reader_lines: Box<dyn Iterator<Item = Result<String, Error>> + Sync + Send>,
    line: Option<Result<String, Arc<Error>>>,
}

impl TokenIterator {
    pub fn from_file(filename: &str) -> std::io::Result<TokenIterator> {
        let file = File::open(filename)?;
        let buf_reader = BufReader::new(file);
        Ok(TokenIterator {
            reader_lines: Box::new(buf_reader.lines()),
            line: None,
        })
    }

    pub fn from_string(string: &str) -> std::io::Result<TokenIterator> {
        let cursor = io::Cursor::new(string.to_owned());
        Ok(TokenIterator {
            reader_lines: Box::new(cursor.lines()),
            line: None,
        })
    }
}

static NEXT_TOKEN: once_cell::sync::Lazy<Regex> =
    Lazy::<Regex>::new(|| Regex::new(r"^\s*(?:(?:\p{L}|['’])+|\p{Punct}|\n|\d+)").unwrap());
static WS: once_cell::sync::Lazy<Regex> = Lazy::<Regex>::new(|| Regex::new(r"^\s+").unwrap());
static NEXT_CHAR: once_cell::sync::Lazy<Regex> = Lazy::<Regex>::new(|| Regex::new(r"^.").unwrap());
pub static ENDS_WITH_ALNUM: once_cell::sync::Lazy<Regex> =
    Lazy::<Regex>::new(|| Regex::new(r"(?:\p{L}|\d)$").unwrap());
pub static STARTS_WITH_ALNUM: once_cell::sync::Lazy<Regex> =
    Lazy::<Regex>::new(|| Regex::new(r"^(?:\p{L}|\d)").unwrap());

// simplistic approach
impl Iterator for TokenIterator {
    type Item = String;

    fn next(&mut self) -> Option<Self::Item> {
        match self.line.as_ref() {
            None => {
                self.line = self.reader_lines.next().map(|x| x.map_err(Arc::new));
                match self.line {
                    None => None,
                    Some(_) => self.next(),
                }
            }
            Some(Err(_)) => None,
            Some(Ok(line)) => {
                if line.is_empty() {
                    self.line = self.reader_lines.next().map(|x| x.map_err(Arc::new));
                    return self.next();
                }

                let m = NEXT_TOKEN.find(line);
                let index = match m {
                    None => {
                        // find unicode char boundary.
                        NEXT_CHAR.find(line)?.end()
                    }
                    Some(m) => m.end(),
                };

                let next = line[..index].to_string();
                self.line = Some(Ok(line[index..].to_string()));
                let next = WS.replace(&next, "");
                if next.len() == 0 {
                    self.next()
                } else {
                    Some(next.into())
                }
            }
        }
    }
}

#[derive(Clone, Debug)]
struct NgramMap<Token> {
    map: Vec<HashMap<Vec<Token>, DiscreteDistribution<Token>>>,
    pre_map: Vec<HashMap<Vec<Token>, HashMap<Token, usize>>>,
}

// #[derive(Debug)]
// enum MarkovError<T> {
//     NotInMap(Vec<T>),
// }

// impl<T> std::fmt::Display for MarkovError<T>
// where
//     T: Debug,
// {
//     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
//         write!(f, "{:?}", self)
//     }
// }

// impl<T> std::error::Error for MarkovError<T> where T: Debug {}

impl<T> NgramMap<T>
where
    T: Hash + Eq + PartialEq + Clone + Debug,
{
    fn new(up_to: usize) -> Self {
        NgramMap {
            pre_map: (1..=up_to).map(|_| HashMap::new()).collect::<Vec<_>>(),
            map: (1..=up_to).map(|_| HashMap::new()).collect::<Vec<_>>(),
        }
    }

    fn query(&self, last_n: &[T]) -> Option<&DiscreteDistribution<T>> {
        // MarkovError<T>> {
        // println!("{:#?}", self.map[last_n.len()]);

        // in case more are given
        let up_to = self.map.len();
        let last_n = &last_n[last_n.len().saturating_sub(up_to)..];

        // That's not very useful. Need to restart ...
        self.map[last_n.len()].get(last_n)
        // .ok_or(MarkovError::NotInMap(last_n.to_vec()))
    }
}

#[derive(Clone, Debug)]
pub struct MarkovChainInternal {
    ngram_map: NgramMap<String>,
    pub state: Vec<String>,
    pub up_to: usize,
}

impl MarkovChainInternal {
    pub fn set_state(&mut self, state: Vec<String>) {
        self.state = state;
    }

    pub fn from_file(filename: &str, up_to: usize) -> Result<MarkovChainInternal, anyhow::Error> {
        let token_iterator = TokenIterator::from_file(filename)?;
        Self::from_token_iterator(token_iterator, up_to)
    }

    pub fn from_token_iterator(
        token_iterator: TokenIterator,
        up_to: usize,
    ) -> Result<MarkovChainInternal, anyhow::Error> {
        let mut ngram_map = NgramMap::new(up_to);

        let mut storage: Storage<String> = Storage::new(up_to);

        for x in token_iterator
            .chain(std::iter::repeat("<EOF>".to_string()).take(up_to - 1))
            .sliding_windows(&mut storage)
        {
            for k in 1..=up_to {
                let mut window = x.iter().cloned().collect::<Vec<String>>();

                // println!("window: {:?}, k = {}", window, k);

                let last = window.splice(k - 1.., vec![]).collect::<Vec<_>>();
                let first_k_minus_1 = window;
                let w = last[0].clone();

                // println!("{:?} {}", first_k_minus_1, w);

                if w != *"<EOF>" {
                    ngram_map.pre_map[k - 1]
                        .entry(first_k_minus_1)
                        .or_insert(Default::default())
                        .entry(w.to_string())
                        .and_modify(|count| *count += 1)
                        .or_insert(1);
                }
            }
        }

        for k in 1..=up_to {
            let mut map_for_k = HashMap::new();

            for next_map in ngram_map.pre_map[k - 1].iter() {
                let sum = next_map.1.values().sum::<usize>() as f64;
                let weights = next_map.1.values().map(|x| (*x as f64) / sum);
                let distro = DiscreteDistribution::from_iterator_and_labels(
                    weights,
                    next_map.1.keys().cloned().collect::<Vec<_>>(),
                );
                map_for_k.insert(next_map.0.to_vec(), distro);
            }

            ngram_map.map.insert(k - 1, map_for_k);
        }

        // println!("{:?}", ngram_map.map[0]);

        Ok(MarkovChainInternal {
            ngram_map,
            state: vec![],
            up_to,
        })
    }
}

impl NextTokenDistributionProvider<String, Vec<String>> for MarkovChainInternal {
    fn reset(
        &mut self,
        init: Vec<String>,
    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send {
        self.state = init;
        ok(())
    }

    fn select(
        &mut self,
        next: String,
    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send {
        // println!("state = {:?}", self.state);

        let range = self.state.len().saturating_sub(1)..;

        self.state = self
            .state
            // .splice(1..=self.up_to - 1, vec![])
            // cmp with len of state?
            .splice(range, vec![])
            // .splice(0..=1, vec![])
            .chain(vec![next])
            .collect();
        ok(())
    }

    fn get(
        &mut self,
    ) -> impl std::future::Future<
        Output = Result<crate::DiscreteDistribution<String>, anyhow::Error>,
    > + Send {
        for k in (0..=self.up_to).rev() {
            let last_n: &[String] = &self.state.as_slice()
                [std::cmp::max(0, self.state.len() as isize - k as isize) as usize..];

            let result = self.ngram_map.query(last_n);

            if let Some(result) = result {
                return ok(result.clone());
            }
        }
        // if properly initialized, should not occur:
        unreachable!()
        // return anyhow::Error::new("no entries".to_string());
        // todo!()
    }
}
